from qlib.contrib.data.handler import Alpha158, check_transform_proc
from qlib.data.dataset.handler import DataHandlerLP

DEFAULT_INFER_PROCESSORS = []
DEFAULT_LEARN_PROCESSORS = [{"class": "DropnaLabel"}]


class EmotionAlpha158(Alpha158):
    def __init__(
            self,
            instruments="csi500",
            start_time=None,
            end_time=None,
            freq="day",
            infer_processors=None,
            learn_processors=None,
            fit_start_time=None,
            fit_end_time=None,
            process_type=DataHandlerLP.PTYPE_A,
            filter_pipe=None,
            inst_processor=None,
            **kwargs
    ):
        if learn_processors is None:
            learn_processors = DEFAULT_LEARN_PROCESSORS
        if infer_processors is None:
            infer_processors = DEFAULT_INFER_PROCESSORS
        infer_processors = check_transform_proc(infer_processors, fit_start_time, fit_end_time)
        learn_processors = check_transform_proc(learn_processors, fit_start_time, fit_end_time)

        data_loader = {
            "class": "QlibDataLoader",
            "kwargs": {
                "config": {
                    "feature": self.get_feature_config(),
                    "label": kwargs.pop("label", self.get_label_config()),
                },
                "filter_pipe": filter_pipe,
                "freq": freq,
                "inst_processor": inst_processor,
            },
        }
        super(Alpha158, self).__init__(
            instruments=instruments,
            start_time=start_time,
            end_time=end_time,
            data_loader=data_loader,
            infer_processors=infer_processors,
            learn_processors=learn_processors,
            process_type=process_type,
            **kwargs
        )

    def get_label_config(self):
        return ["Ref($close, -5)/Ref($close, -1) - 1"], ["LABEL0"]

    def get_feature_config(self):
        conf = {
            "kbar": {},
            "price": {
                "windows": [0],
                "feature": ["OPEN", "HIGH", "LOW", "VWAP"],
            },
            "rolling": {},
        }
        fields_emotion = [f"$emotion_{i + 1}" for i in range(15)]
        names_emotion = [f"emotion_{i + 1}" for i in range(15)]
        fields, names = self.parse_config_to_fields(conf)
        return fields_emotion + fields, names_emotion + names
